#coding = 'utf-8'

import rospy
from mavros_msgs.srv import SetMode, CommandBool, CommandTOL
from mavros_msgs.msg import State
from std_msgs.msg import String
from PyQt5 import QtWidgets
import sys
from UI_px4 import Ui_Form 

class PX4:
    def __init__(self, uav_id, scroll_area):
        """
        初始化 PX4 对象，连接到指定的无人机。
        :param uav_id: 无人机的编号，例如 "uav0"。
        """
        self.uav_id = uav_id
        self.uav_name = f"uav{self.uav_id}"
        self.namespace = f"/{self.uav_name}/mavros"
        print(f"drone name: {self.uav_name}")
        print(f"drone namespace: {self.namespace}")
        self.state = State()
        self.scroll_area = scroll_area

        # 初始化 UI
        self.ui = Ui_Form()
        self.form_widget = QtWidgets.QWidget()  # 创建一个 QWidget 作为表单的容器
        self.ui.setupUi(self.form_widget)

        # 将表单容器设置为滚动区域的内容部件
        self.scroll_area.setWidget(self.form_widget)


        # 初始化 ROS 节点
        rospy.init_node(f"px4_control_{self.uav_id}")
        # 订阅无人机状态
        rospy.Subscriber(f"{self.namespace}/mavros/state", State, self.state_callback)

        # 等待 ROS 服务可用
        rospy.wait_for_service(f"{self.namespace}/set_mode")
        rospy.wait_for_service(f"{self.namespace}/cmd/arming")
        rospy.wait_for_service(f"{self.namespace}/cmd/takeoff")

        self.set_mode_client = rospy.ServiceProxy(f"{self.namespace}/set_mode", SetMode)
        self.arming_client = rospy.ServiceProxy(f"{self.namespace}/cmd/arming", CommandBool)
        self.takeoff_client = rospy.ServiceProxy(f"{self.namespace}/cmd/takeoff", CommandTOL)

    def state_callback(self, data):
        """
        回调函数，更新无人机状态。
        """
        # print("====>state: ", data.connected)
        self.state = data
        status = "Online" if self.state.connected else "Offline"
        # self.ui.lineEdit.setText()
        
    def set_mode(self, mode):
        """
        设置无人机的飞行模式。
        :param mode: 目标模式，例如 "OFFBOARD" 或 "AUTO.TAKEOFF"。
        :return: 设置成功返回 True，否则返回 False。
        """
        if self.state.mode == mode:
            rospy.loginfo(f"Mode already set to {mode}")
            return True

        try:
            result = self.set_mode_client(base_mode=0, custom_mode=mode)
            if result.mode_sent:
                rospy.loginfo(f"Mode set to {mode}")
                return True
            else:
                rospy.logerr(f"Failed to set mode to {mode}")
                return False
        except rospy.ServiceException as e:
            rospy.logerr(f"Service call failed: {e}")
            return False

    def arm(self):
        """
        解锁无人机。
        :return: 解锁成功返回 True，否则返回 False。
        """
        try:
            result = self.arming_client(value=True)
            if result.success:
                rospy.loginfo("Vehicle armed")
                return True
            else:
                rospy.logerr("Failed to arm vehicle")
                return False
        except rospy.ServiceException as e:
            rospy.logerr(f"Service call failed: {e}")
            return False

    def takeoff(self, altitude=2.0):
        """
        执行起飞操作。
        :param altitude: 起飞高度，默认为 2.0 米。
        :return: 起飞成功返回 True，否则返回 False。
        """
        try:
          if self.set_mode("AUTO.TAKEOFF"):
              rospy.sleep(2)  # 等待模式切换完成

              # 解锁无人机
              if self.arm():
                  rospy.sleep(2)  # 等待解锁完成
        except rospy.ServiceException as e:
            rospy.logerr(f"Service call failed: {e}")
            return False

    def is_armed(self):
        """
        检查无人机是否已解锁。
        :return: 如果无人机已解锁，返回 True，否则返回 False。
        """
        return self.state.armed

    def is_connected(self):
        """
        检查无人机是否已连接。
        :return: 如果无人机已连接，返回 True，否则返回 False。
        """
        return self.state.connected
    